package core

import (
	"crypto/sha256"
	"github.com/hashicorp/go-hclog"
	"sync"
)

// RBC implements Reliable Broadcast protocol
type RBC struct {
	node   *Node
	logger hclog.Logger
	
	// Echo and Ready message storage
	echoMsgs  map[int]map[int][]byte // [SN][Sender]Data
	readyMsgs map[int]map[int][]byte // [SN][Sender]Hash
	
	// Output tracking
	outputs   map[int][]byte // [SN]Data
	delivered map[int]bool   // [SN]bool
	
	// Threshold values
	t int // threshold = f+1
	n int // total nodes
	
	mutex sync.RWMutex
}

// NewRBC creates a new RBC instance
func NewRBC(node *Node) *RBC {
	return &RBC{
		node:      node,
		logger:    hclog.New(&hclog.LoggerOptions{Name: "dumbo-rbc", Level: hclog.Level(node.Config.LogLevel)}),
		echoMsgs:  make(map[int]map[int][]byte),
		readyMsgs: make(map[int]map[int][]byte),
		outputs:   make(map[int][]byte),
		delivered: make(map[int]bool),
		t:         node.F + 1,
		n:         node.N,
	}
}

// Broadcast initiates RBC for a given sequence number and data
func (r *RBC) Broadcast(sn int, data []byte) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	
	if r.delivered[sn] {
		return // Already delivered
	}
	
	// 开始RBC延迟计时
	r.node.delayMonitor.StartRBC(sn)
	
	// Hash the data
	hash := sha256.Sum256(data)
	
	// Send Echo message to all nodes
	echoMsg := RBCEchoMsg{
		SN:     sn,
		Sender: r.node.Id,
		Data:   data,
		Hash:   hash[:],
	}
	
	r.logger.Debug("RBC broadcasting", "sn", sn, "data_len", len(data))
	
	// Broadcast echo message
	r.node.PlainBroadcast(RBCEchoMsgTag, echoMsg, nil)
}

// HandleEchoMsg processes received echo messages
func (r *RBC) HandleEchoMsg(msg *RBCEchoMsg) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	
	sn := msg.SN
	sender := msg.Sender
	
	// Initialize maps if needed
	if r.echoMsgs[sn] == nil {
		r.echoMsgs[sn] = make(map[int][]byte)
	}
	
	// Store echo message
	r.echoMsgs[sn][sender] = msg.Data
	
	r.logger.Debug("RBC received echo", "sn", sn, "sender", sender, "echo_count", len(r.echoMsgs[sn]))
	
	// Check if we have t echo messages with same data
	if len(r.echoMsgs[sn]) >= r.t {
		// Find the most common data
		dataCount := make(map[string]int)
		for _, data := range r.echoMsgs[sn] {
			dataCount[string(data)]++
		}
		
		// Find data with at least t occurrences
		var commonData []byte
		for data, count := range dataCount {
			if count >= r.t {
				commonData = []byte(data)
				break
			}
		}
		
		if commonData != nil {
			// Send Ready message
			hash := sha256.Sum256(commonData)
			readyMsg := RBCReadyMsg{
				SN:     sn,
				Sender: r.node.Id,
				Hash:   hash[:],
			}
			
			r.logger.Debug("RBC sending ready", "sn", sn)
			r.node.PlainBroadcast(RBCReadyMsgTag, readyMsg, nil)
		}
	}
}

// HandleReadyMsg processes received ready messages
func (r *RBC) HandleReadyMsg(msg *RBCReadyMsg) {
	r.mutex.Lock()
	defer r.mutex.Unlock()
	
	sn := msg.SN
	sender := msg.Sender
	
	// Initialize maps if needed
	if r.readyMsgs[sn] == nil {
		r.readyMsgs[sn] = make(map[int][]byte)
	}
	
	// Store ready message
	r.readyMsgs[sn][sender] = msg.Hash
	
	r.logger.Debug("RBC received ready", "sn", sn, "sender", sender, "ready_count", len(r.readyMsgs[sn]))
	
	// Check if we have t ready messages with same hash
	if len(r.readyMsgs[sn]) >= r.t {
		// Find the most common hash
		hashCount := make(map[string]int)
		for _, hash := range r.readyMsgs[sn] {
			hashCount[string(hash)]++
		}
		
		// Find hash with at least t occurrences
		var commonHash []byte
		for hash, count := range hashCount {
			if count >= r.t {
				commonHash = []byte(hash)
				break
			}
		}
		
		if commonHash != nil && !r.delivered[sn] {
			// Try to find the corresponding data
			if r.echoMsgs[sn] != nil {
				for _, data := range r.echoMsgs[sn] {
					dataHash := sha256.Sum256(data)
					if string(dataHash[:]) == string(commonHash) {
						// Deliver the data
						r.deliver(sn, data)
						break
					}
				}
			}
		}
	}
	
	// Check if we have 2t ready messages
	if len(r.readyMsgs[sn]) >= 2*r.t {
		// Find the most common hash
		hashCount := make(map[string]int)
		for _, hash := range r.readyMsgs[sn] {
			hashCount[string(hash)]++
		}
		
		// Find hash with at least 2t occurrences
		var commonHash []byte
		for hash, count := range hashCount {
			if count >= 2*r.t {
				commonHash = []byte(hash)
				break
			}
		}
		
		if commonHash != nil && !r.delivered[sn] {
			// Try to find the corresponding data
			if r.echoMsgs[sn] != nil {
				for _, data := range r.echoMsgs[sn] {
					dataHash := sha256.Sum256(data)
					if string(dataHash[:]) == string(commonHash) {
						// Deliver the data
						r.deliver(sn, data)
						break
					}
				}
			}
		}
	}
}

// deliver outputs the data for a given sequence number
func (r *RBC) deliver(sn int, data []byte) {
	if r.delivered[sn] {
		return
	}
	
	r.delivered[sn] = true
	r.outputs[sn] = data
	
	// 结束RBC延迟计时并记录
	delay := r.node.delayMonitor.EndRBC(sn)
	r.logger.Info("RBC delivered", "sn", sn, "data_len", len(data), "delay_ms", delay.Milliseconds())
	
	// Send output message
	outputMsg := RBCOutputMsg{
		SN:     sn,
		Sender: r.node.Id,
		Data:   data,
	}
	
	r.node.PlainBroadcast(RBCOutputMsgTag, outputMsg, nil)
}

// GetOutput returns the output for a given sequence number
func (r *RBC) GetOutput(sn int) ([]byte, bool) {
	r.mutex.RLock()
	defer r.mutex.RUnlock()
	
	data, exists := r.outputs[sn]
	return data, exists
}

// IsDelivered checks if a sequence number has been delivered
func (r *RBC) IsDelivered(sn int) bool {
	r.mutex.RLock()
	defer r.mutex.RUnlock()
	
	return r.delivered[sn]
}
